#  -*- coding: utf-8 -*-

"""
普量学院量化投资课程系列案例源码包
普量学院版权所有
仅用于教学目的，严禁转发和用于盈利目的，违者必究
©Plouto-Quants All Rights Reserved

普量学院助教微信：niuxiaomi3
"""

"""
MongoDB索引管理脚本
为daily、daily_hfq、basic、finance_report等表创建最优索引
提高数据新增和查询效率
"""

import time
from pymongo import ASCENDING, DESCENDING
from database import DB_CONN

def create_daily_indexes():
    """
    为daily表创建索引
    """
    print("🔧 开始为daily表创建索引...")
    
    collection = DB_CONN['daily']
    
    # 1. 主键索引：code + date + index (最重要的索引)
    # 这是upsert操作的查询条件，需要唯一性
    print("创建主键索引: code + date + index")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('code', ASCENDING),
            ('date', ASCENDING),
            ('index', ASCENDING)
        ], 
        unique=True,  # 唯一索引，防止重复数据
        background=True,  # 后台创建，不阻塞其他操作
        name='idx_code_date_index'
        )
        print(f"✅ 主键索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 主键索引创建失败: {e}")
    
    # 2. 股票代码索引：code
    # 用于按股票查询历史数据
    print("创建股票代码索引: code")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('code', ASCENDING)
        ], 
        background=True,
        name='idx_code'
        )
        print(f"✅ 股票代码索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 股票代码索引创建失败: {e}")
    
    # 3. 日期索引：date
    # 用于按日期范围查询
    print("创建日期索引: date")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('date', ASCENDING)
        ], 
        background=True,
        name='idx_date'
        )
        print(f"✅ 日期索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 日期索引创建失败: {e}")
    
    # 4. 指数标识索引：index
    # 用于区分股票和指数
    print("创建指数标识索引: index")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('index', ASCENDING)
        ], 
        background=True,
        name='idx_index'
        )
        print(f"✅ 指数标识索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 指数标识索引创建失败: {e}")
    
    # 5. 交易状态索引：is_trading
    # 用于查询正常交易的股票
    print("创建交易状态索引: is_trading")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('is_trading', ASCENDING)
        ], 
        background=True,
        name='idx_is_trading'
        )
        print(f"✅ 交易状态索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 交易状态索引创建失败: {e}")
    
    # 6. PE值索引：pe
    # 用于PE值范围查询和排序
    print("创建PE值索引: pe")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('pe', ASCENDING)
        ], 
        background=True,
        name='idx_pe'
        )
        print(f"✅ PE值索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ PE值索引创建失败: {e}")
    
    # 7. 复合索引：date + pe + is_trading (用于股票池选股)
    # 这是股票池策略中最常用的查询条件
    print("创建复合索引: date + pe + is_trading")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('date', ASCENDING),
            ('pe', ASCENDING),
            ('is_trading', ASCENDING)
        ], 
        background=True,
        name='idx_date_pe_trading'
        )
        print(f"✅ 复合索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 复合索引创建失败: {e}")
    
    # 8. 复合索引：code + date (用于时间序列查询)
    # 用于获取某只股票的历史数据
    print("创建复合索引: code + date")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('code', ASCENDING),
            ('date', ASCENDING)
        ], 
        background=True,
        name='idx_code_date'
        )
        print(f"✅ 复合索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 复合索引创建失败: {e}")
    
    print("🎉 daily表索引创建完成！\n")


def create_daily_hfq_indexes():
    """
    为daily_hfq表创建索引（与daily表相同的索引策略）
    """
    print("🔧 开始为daily_hfq表创建索引...")
    
    collection = DB_CONN['daily_hfq']
    
    # 索引策略与daily表相同
    indexes = [
        {
            'fields': [('code', ASCENDING), ('date', ASCENDING), ('index', ASCENDING)],
            'options': {'unique': True, 'background': True, 'name': 'idx_code_date_index'},
            'description': '主键索引: code + date + index'
        },
        {
            'fields': [('code', ASCENDING)],
            'options': {'background': True, 'name': 'idx_code'},
            'description': '股票代码索引: code'
        },
        {
            'fields': [('date', ASCENDING)],
            'options': {'background': True, 'name': 'idx_date'},
            'description': '日期索引: date'
        },
        {
            'fields': [('index', ASCENDING)],
            'options': {'background': True, 'name': 'idx_index'},
            'description': '指数标识索引: index'
        },
        {
            'fields': [('is_trading', ASCENDING)],
            'options': {'background': True, 'name': 'idx_is_trading'},
            'description': '交易状态索引: is_trading'
        },
        {
            'fields': [('code', ASCENDING), ('date', ASCENDING)],
            'options': {'background': True, 'name': 'idx_code_date'},
            'description': '复合索引: code + date'
        }
    ]
    
    for index_info in indexes:
        print(f"创建{index_info['description']}")
        start_time = time.time()
        try:
            result = collection.create_index(index_info['fields'], **index_info['options'])
            print(f"✅ {index_info['description']}创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
        except Exception as e:
            print(f"❌ {index_info['description']}创建失败: {e}")
    
    print("🎉 daily_hfq表索引创建完成！\n")


def create_basic_indexes():
    """
    为basic表创建索引
    """
    print("🔧 开始为basic表创建索引...")
    
    collection = DB_CONN['basic']
    
    # 1. 主键索引：code + date
    print("创建主键索引: code + date")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('code', ASCENDING),
            ('date', ASCENDING)
        ], 
        unique=True,
        background=True,
        name='idx_code_date'
        )
        print(f"✅ 主键索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 主键索引创建失败: {e}")
    
    # 2. 日期索引：date
    print("创建日期索引: date")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('date', ASCENDING)
        ], 
        background=True,
        name='idx_date'
        )
        print(f"✅ 日期索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 日期索引创建失败: {e}")
    
    print("🎉 basic表索引创建完成！\n")


def create_finance_report_indexes():
    """
    为finance_report表创建索引
    """
    print("🔧 开始为finance_report表创建索引...")
    
    collection = DB_CONN['finance_report']
    
    # 1. 主键索引：code + report_date
    print("创建主键索引: code + report_date")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('code', ASCENDING),
            ('report_date', ASCENDING)
        ], 
        unique=True,
        background=True,
        name='idx_code_report_date'
        )
        print(f"✅ 主键索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 主键索引创建失败: {e}")
    
    # 2. 复合索引：code + announced_date (用于PE计算)
    print("创建复合索引: code + announced_date")
    start_time = time.time()
    try:
        result = collection.create_index([
            ('code', ASCENDING),
            ('announced_date', DESCENDING)  # 降序，用于获取最新公告
        ], 
        background=True,
        name='idx_code_announced_date'
        )
        print(f"✅ 复合索引创建完成: {result}, 耗时: {time.time() - start_time:.2f}秒")
    except Exception as e:
        print(f"❌ 复合索引创建失败: {e}")
    
    print("🎉 finance_report表索引创建完成！\n")


def drop_all_indexes():
    """
    删除所有索引（保留_id索引）
    """
    print("🗑️  开始删除所有索引...")
    
    collections = ['daily', 'daily_hfq', 'basic', 'finance_report']
    
    for collection_name in collections:
        try:
            collection = DB_CONN[collection_name]
            result = collection.drop_indexes()
            print(f"✅ {collection_name}表索引删除完成")
        except Exception as e:
            print(f"❌ {collection_name}表索引删除失败: {e}")
    
    print("🎉 所有索引删除完成！\n")


def show_all_indexes():
    """
    显示所有表的索引信息
    """
    print("📋 显示所有表的索引信息...")
    
    collections = ['daily', 'daily_hfq', 'basic', 'finance_report']
    
    for collection_name in collections:
        try:
            collection = DB_CONN[collection_name]
            indexes = collection.list_indexes()
            
            print(f"\n📊 {collection_name}表的索引:")
            print("-" * 50)
            
            for index in indexes:
                name = index.get('name', 'unknown')
                key = index.get('key', {})
                unique = index.get('unique', False)
                
                print(f"  索引名: {name}")
                print(f"  字段: {dict(key)}")
                print(f"  唯一性: {'是' if unique else '否'}")
                print("-" * 30)
                
        except Exception as e:
            print(f"❌ 获取{collection_name}表索引信息失败: {e}")
    
    print("\n🎉 索引信息显示完成！")


def analyze_index_performance():
    """
    分析索引性能
    """
    print("🔍 开始分析索引性能...")
    
    # 测试查询性能
    test_queries = [
        {
            'collection': 'daily',
            'query': {'code': '000001', 'date': '2015-01-01'},
            'description': '单股票单日期查询'
        },
        {
            'collection': 'daily',
            'query': {'date': {'$gte': '2015-01-01', '$lte': '2015-12-31'}},
            'description': '日期范围查询'
        },
        {
            'collection': 'daily',
            'query': {'code': '000001'},
            'description': '单股票历史数据查询'
        },
        {
            'collection': 'daily',
            'query': {'date': '2015-01-01', 'pe': {'$gt': 0, '$lt': 30}, 'is_trading': True},
            'description': '股票池选股查询'
        }
    ]
    
    for test in test_queries:
        try:
            collection = DB_CONN[test['collection']]
            
            # 执行查询并测量时间
            start_time = time.time()
            cursor = collection.find(test['query']).limit(100)
            results = list(cursor)
            query_time = time.time() - start_time
            
            # 获取查询计划
            explain = collection.find(test['query']).explain()
            
            print(f"\n📈 {test['description']}:")
            print(f"  查询时间: {query_time:.4f}秒")
            print(f"  结果数量: {len(results)}条")
            print(f"  扫描文档数: {explain.get('executionStats', {}).get('totalDocsExamined', 'N/A')}")
            print(f"  使用索引: {explain.get('executionStats', {}).get('executionSuccess', 'N/A')}")
            
        except Exception as e:
            print(f"❌ 分析{test['description']}失败: {e}")
    
    print("\n🎉 索引性能分析完成！")


def main():
    """
    主函数
    """
    print("=" * 60)
    print("🚀 MongoDB索引管理工具")
    print("=" * 60)
    
    while True:
        print("\n请选择操作:")
        print("1. 创建所有索引")
        print("2. 创建daily表索引")
        print("3. 创建daily_hfq表索引")
        print("4. 创建basic表索引")
        print("5. 创建finance_report表索引")
        print("6. 删除所有索引")
        print("7. 显示所有索引")
        print("8. 分析索引性能")
        print("9. 退出")
        
        choice = input("\n请输入选择 (1-9): ")
        
        if choice == '1':
            total_start_time = time.time()
            create_daily_indexes()
            create_daily_hfq_indexes()
            create_basic_indexes()
            create_finance_report_indexes()
            total_time = time.time() - total_start_time
            print(f"🎉 所有索引创建完成！总耗时: {total_time:.2f}秒")
            
        elif choice == '2':
            create_daily_indexes()
            
        elif choice == '3':
            create_daily_hfq_indexes()
            
        elif choice == '4':
            create_basic_indexes()
            
        elif choice == '5':
            create_finance_report_indexes()
            
        elif choice == '6':
            drop_all_indexes()
            
        elif choice == '7':
            show_all_indexes()
            
        elif choice == '8':
            analyze_index_performance()
            
        elif choice == '9':
            print("👋 再见！")
            break
            
        else:
            print("❌ 无效选择，请重新输入！")


if __name__ == "__main__":
    main() 